Skip to content

fix(flux): support .weight variant for attention norm parameter mappings#9190

Open
Fortyseven wants to merge 3 commits into
invoke-ai:mainfrom
Fortyseven:flux-norm-weight-mapping
Open

fix(flux): support .weight variant for attention norm parameter mappings#9190
Fortyseven wants to merge 3 commits into
invoke-ai:mainfrom
Fortyseven:flux-norm-weight-mapping

Conversation

@Fortyseven
Copy link
Copy Markdown

@Fortyseven Fortyseven commented May 16, 2026

Summary

Some Flux model checkpoints use *.weight instead of *.scale for attention normalization layer parameters (query_norm and key_norm). The model loader's key mapping functions only checked for .scale, causing these checkpoints to fail loading with unmapped keys.

This change updates all affected mapping functions across multiple Flux model type handlers to also match the .weight variant, mapping it to the same internal parameter name as .scale.

Related Issues / Discussions

https://discord.com/channels/1020123559063990373/1083864753543331981/1504216036168171570

QA Instructions

This was tested with Flux2-Klein-9B-True-V2. (Q4_K specifically, but other sizes exhibit the issue too.)

Installation goes fine, and it's recognized. However, attempting a generation results in:

Server Error

RuntimeError: Error(s) in loading state_dict for Flux2Transformer2DModel: Missing key(s) in state_dict: "transformer_blocks.0.attn.norm_q.weight", "transformer_blocks.0.attn.norm_k.weight",...

After this patch: image generation works.

Note

It's important to note that while I have quite a long history in software development, I am on the very fringes of this particular tech. So I may be overlooking something in this 'fix' that's just outside my understanding. ;)

@github-actions github-actions Bot added python PRs that change python files backend PRs that change backend files labels May 16, 2026
@lstein lstein added the v6.13.x label May 17, 2026
@lstein lstein moved this to 6.13.5 LIBRARY UPDATES in Invoke - Community Roadmap May 17, 2026
@lstein lstein added 6.13.5 Library Updates and removed v6.13.x labels May 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

6.13.5 Library Updates backend PRs that change backend files python PRs that change python files

Projects

Status: 6.13.5 LIBRARY UPDATES

Development

Successfully merging this pull request may close these issues.

3 participants